import time

from collections import deque
import numpy as np
import mujoco
import mujoco.viewer
import torch
import onnxruntime as ort
import yaml


def get_gravity_orientation(quaternion):
    qw = quaternion[0]
    qx = quaternion[1]
    qy = quaternion[2]
    qz = quaternion[3]

    gravity_orientation = np.zeros(3)

    gravity_orientation[0] = 2 * (-qz * qx + qw * qy)
    gravity_orientation[1] = -2 * (qz * qy + qw * qx)
    gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz)

    return gravity_orientation


def pd_control(target_q, q, kp, target_dq, dq, kd):
    return (target_q - q) * kp + (target_dq - dq) * kd


if __name__ == '__main__':
    with open(f"./g1_config/g1.yaml", 'r') as file:
        config = yaml.safe_load(file)

        kps = np.array(config["control"]["stiffness"], dtype=np.float32)
        kds = np.array(config["control"]["damping"], dtype=np.float32)
        torque_limit = np.array(config["control"]["torque_limit"])

        default_joint_angles = np.array(config["default_joint_angles"], dtype=np.float32)

    action = np.zeros(config["actions_dim"], dtype=np.float32)
    target_dof_pos = default_joint_angles.copy()
    obs = np.zeros([1, config["observation_dim"]], dtype=np.float32)
    hist_obs = np.zeros([1, (config["frame_stack"] + 1) * config["observation_dim"]], dtype=np.float32)

    counter = 0

    policy = ort.InferenceSession(config["policy_path"])

    mj_model = mujoco.MjModel.from_xml_path(config["mjcf_path"])
    mj_data = mujoco.MjData(mj_model)
    mj_model.opt.timestep = config["simulation_dt"]

    initial_qpos = mj_data.qpos.copy()
    initial_qpos[7:] = default_joint_angles
    mujoco.mj_setState(mj_model, mj_data, initial_qpos, mujoco.mjtState.mjSTATE_QPOS)
    mujoco.mj_forward(mj_model, mj_data)

    with mujoco.viewer.launch_passive(mj_model, mj_data) as viewer:
        while viewer.is_running() and mj_data.time < config["simulation_duration"]:
            step_start = time.time()
            tau = pd_control(target_dof_pos, mj_data.qpos[7:], kps, np.zeros_like(kds), mj_data.qvel[6:], kds)
            tau = np.clip(tau, -torque_limit, torque_limit)
            mj_data.ctrl[:] = tau
            mujoco.mj_step(mj_model, mj_data)

            counter += 1
            if counter % config["control_decimation"] == 0:
                qj = mj_data.qpos[7:]
                dqj = mj_data.qvel[6:]
                quat = mj_data.qpos[3:7]
                omega = mj_data.qvel[3:6]

                gravity_orientation = get_gravity_orientation(quat)

                ref_motion_phase = (counter + 1) * config["simulation_dt"] / config["cycle_time"]
                ref_motion_phase = np.clip(ref_motion_phase, 0, 1)

                obs[0, :23] = action * config["obs_scales"]["actions"]
                obs[0, 23:26] = omega * config["obs_scales"]["base_ang_vel"]
                obs[0, 26:49] = (qj - default_joint_angles) * config["obs_scales"]["dof_pos"]
                obs[0, 49:72] = dqj * config["obs_scales"]["dof_vel"]
                obs[0, 72:75] = gravity_orientation * config["obs_scales"]["projected_gravity"]
                obs[0, 75] = ref_motion_phase * config["obs_scales"]["ref_motion_phase"]

                hist_obs[0, 72 + 23: 72 + 23 * 4] = hist_obs[0, 72: 72 + 23 * 3] * config["obs_scales"]["history_actor"]
                hist_obs[0, 164 + 3: 164 + 3 * 4] = hist_obs[0, 164: 164 + 3 * 3] * config["obs_scales"][
                    "history_actor"]
                hist_obs[0, 176 + 23: 176 + 23 * 4] = hist_obs[0, 176: 176 + 23 * 3] * config["obs_scales"][
                    "history_actor"]
                hist_obs[0, 268 + 23: 268 + 23 * 4] = hist_obs[0, 268: 268 + 23 * 3] * config["obs_scales"][
                    "history_actor"]
                hist_obs[0, 360 + 3: 360 + 3 * 4] = hist_obs[0, 360: 360 + 3 * 3] * config["obs_scales"][
                    "history_actor"]
                hist_obs[0, 372 + 1: 372 + 1 * 4] = hist_obs[0, 372: 372 + 1 * 3] * config["obs_scales"][
                    "history_actor"]

                hist_obs[0, 72: 72 + 23] = hist_obs[0, :23] * config["obs_scales"]["history_actor"]
                hist_obs[0, 164: 164 + 3] = hist_obs[0, 23: 26] * config["obs_scales"]["history_actor"]
                hist_obs[0, 176: 176 + 23] = hist_obs[0, 26: 49] * config["obs_scales"]["history_actor"]
                hist_obs[0, 268: 268 + 23] = hist_obs[0, 49: 72] * config["obs_scales"]["history_actor"]
                hist_obs[0, 360: 360 + 3] = hist_obs[0, 376: 379] * config["obs_scales"]["history_actor"]
                hist_obs[0, 372: 372 + 1] = hist_obs[0, 379] * config["obs_scales"]["history_actor"]

                hist_obs[0, :72] = obs[0, :72]
                hist_obs[0, 376: 379] = obs[0, 72:75]
                hist_obs[0, 379] = obs[0, 75]

                hist_obs = np.clip(hist_obs, -config["normalization"]["clip_observations"],
                                   config["normalization"]["clip_observations"])

                policy_input = {policy.get_inputs()[0].name: hist_obs}
                action = policy.run([policy.get_outputs()[0].name], policy_input)[0]
                action = np.clip(action, -config["normalization"]["clip_actions"],
                                 config["normalization"]["clip_actions"])

                target_dof_pos = action * config["control"]["action_scale"] + default_joint_angles
            viewer.sync()

            time_until_next_step = mj_model.opt.timestep - (time.time() - step_start)
            if time_until_next_step > 0:
                time.sleep(time_until_next_step)
